{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# DiscreteConditional"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\nAtlanta, Georgia 30332-0415\nAll Rights Reserved\n\nAuthors: Frank Dellaert, et al. (see THANKS for the full author list)\n\nSee LICENSE for the license information"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/discrete/doc/DiscreteConditional.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `DiscreteConditional` represents a conditional probability distribution, $P(\\text{Frontal} | \\text{Parents})$, for a set of discrete variables. It is a fundamental building block for representing directed graphical models, like Bayes Nets.\n",
        "\n",
        "It is a `DecisionTreeFactor` underneath, but with the additional structure that it distinguishes between frontal (child) and parent variables. The stored values are the probabilities, and for any given assignment to the parent variables, the probabilities of the frontal variable(s) sum to 1."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "import graphviz\n",
        "\n",
        "from gtsam.symbol_shorthand import C, S, R"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a DiscreteConditional\n",
        "\n",
        "A `DiscreteConditional` is created by specifying the frontal (child) key(s), the parent key(s), and the conditional probability table (CPT). The CPT can be given as a list of numbers or a formatted string."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- P(Cloudy) ---\n",
            "Discrete Conditional\n",
            " P( c0 ):\n",
            " Leaf  0.5\n",
            "\n",
            "\n",
            "--- P(Sprinkler | Cloudy) ---\n",
            "Discrete Conditional\n",
            " P( s0 | c0 ):\n",
            " Choice(s0) \n",
            " 0 Choice(c0) \n",
            " 0 0 Leaf  0.5\n",
            " 0 1 Leaf  0.9\n",
            " 1 Choice(c0) \n",
            " 1 0 Leaf  0.5\n",
            " 1 1 Leaf  0.1\n",
            "\n",
            "\n",
            "--- P(Rain | Cloudy) ---\n",
            "Discrete Conditional\n",
            " P( r0 | c0 ):\n",
            " Choice(r0) \n",
            " 0 Choice(c0) \n",
            " 0 0 Leaf  0.2\n",
            " 0 1 Leaf  0.8\n",
            " 1 Choice(c0) \n",
            " 1 0 Leaf  0.8\n",
            " 1 1 Leaf  0.2\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Define keys for three binary variables: Cloudy, Sprinkler, Rain\n",
        "# The first element is the gtsam Key, the second is the cardinality.\n",
        "Cloudy = (C(0), 2)\n",
        "Sprinkler = (S(0), 2)\n",
        "Rain = (R(0), 2)\n",
        "\n",
        "# Create P(Cloudy), a conditional with no parents (a prior).\n",
        "# This is equivalent to a DiscreteDistribution.\n",
        "p_C = gtsam.DiscreteConditional(Cloudy, [], \"0.5/0.5\")\n",
        "print(\"--- P(Cloudy) ---\")\n",
        "p_C.print()\n",
        "\n",
        "# Create P(Sprinkler | Cloudy)\n",
        "# CPT is ordered by parent assignments: C=0, C=1\n",
        "# For C=0 (false), P(S=0|C=0)=0.5, P(S=1|C=0)=0.5\n",
        "# For C=1 (true),  P(S=0|C=1)=0.9, P(S=1|C=1)=0.1\n",
        "p_S_given_C = gtsam.DiscreteConditional(Sprinkler, [Cloudy], \"0.5/0.5 0.9/0.1\")\n",
        "print(\"\\n--- P(Sprinkler | Cloudy) ---\")\n",
        "p_S_given_C.print()\n",
        "\n",
        "# Create P(Rain | Cloudy)\n",
        "# For C=0 (false), P(R=0|C=0)=0.2, P(R=1|C=0)=0.8\n",
        "# For C=1 (true),  P(R=0|C=1)=0.8, P(R=1|C=1)=0.2\n",
        "p_R_given_C = gtsam.DiscreteConditional(Rain, [Cloudy], \"0.2/0.8 0.8/0.2\")\n",
        "print(\"\\n--- P(Rain | Cloudy) ---\")\n",
        "p_R_given_C.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Operations on DiscreteConditional"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "P(S=0|C=1) = 0.9 (log: -0.10536051565782628)\n",
            "\n",
            "Sample for Sprinkler given Cloudy=true: 0\n",
            "Most likely state for Sprinkler given Cloudy=true: 0\n"
          ]
        }
      ],
      "source": [
        "# --- Evaluation ---\n",
        "# Evaluate the probability for a full assignment of variables.\n",
        "values = gtsam.DiscreteValues()\n",
        "values[C(0)] = 1 # Cloudy = true\n",
        "values[S(0)] = 0 # Sprinkler = false\n",
        "\n",
        "prob = p_S_given_C.evaluate(values)\n",
        "log_prob = p_S_given_C.logProbability(values)\n",
        "print(f\"P(S=0|C=1) = {prob} (log: {log_prob})\")\n",
        "\n",
        "# --- Sampling ---\n",
        "# Sample the frontal variable given an assignment for the parents.\n",
        "parent_values = gtsam.DiscreteValues()\n",
        "parent_values[C(0)] = 1 # Condition on Cloudy = true\n",
        "sampled_sprinkler = p_S_given_C.sample(parent_values)\n",
        "print(f\"\\nSample for Sprinkler given Cloudy=true: {sampled_sprinkler}\")\n",
        "\n",
        "# --- Argmax (Most Probable Explanation) ---\n",
        "# Find the most likely assignment of the frontal variable given parents.\n",
        "mpe_sprinkler = p_S_given_C.argmax(parent_values)\n",
        "print(f\"Most likely state for Sprinkler given Cloudy=true: {mpe_sprinkler}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "P(S | C=false):\n",
            "Discrete Conditional\n",
            " P( s0 ):\n",
            " Leaf  0.5\n",
            "\n",
            "\n",
            "Likelihood L(C | S=true):\n",
            "DecisionTreeFactor\n",
            " f[ (c0,2), ]\n",
            " Choice(c0) \n",
            " 0 Leaf  0.5\n",
            " 1 Leaf  0.1\n"
          ]
        }
      ],
      "source": [
        "# --- Choose (Conditioning on parent values) ---\n",
        "# Restricting a conditional on a parent value yields a new conditional\n",
        "# (or a prior if all parents are specified).\n",
        "\n",
        "# Let's fix Cloudy=false (0)\n",
        "given_C_false = gtsam.DiscreteValues()\n",
        "given_C_false[C(0)] = 0\n",
        "p_S = p_S_given_C.choose(given_C_false)\n",
        "\n",
        "print(\"P(S | C=false):\")\n",
        "p_S.print()\n",
        "\n",
        "# --- Likelihood ---\n",
        "# Create a likelihood factor on the parents given a value for the child.\n",
        "frontal_values = gtsam.DiscreteValues()\n",
        "frontal_values[S(0)] = 1 # Evidence: Sprinkler=true\n",
        "likelihood_of_C = p_S_given_C.likelihood(frontal_values)\n",
        "\n",
        "print(\"\\nLikelihood L(C | S=true):\")\n",
        "likelihood_of_C.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Displaying with Markdown/HTML\n",
        "The rich display for a conditional shows the full CPT."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<p>  <i>P(s0|c0):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>c0</i></th><th>0</th><th>1</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><td>0.5</td><td>0.5</td></tr>\n",
              "    <tr><th>1</th><td>0.9</td><td>0.1</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/markdown": [
              " *P(s0|c0):*\n",
              "\n",
              "|*c0*|0|1|\n",
              "|:-:|:-:|:-:|\n",
              "|0|0.5|0.5|\n",
              "|1|0.9|0.1|\n"
            ],
            "text/plain": [
              "Discrete Conditional\n",
              " P( s0 | c0 ):\n",
              " Choice(s0) \n",
              " 0 Choice(c0) \n",
              " 0 0 Leaf  0.5\n",
              " 0 1 Leaf  0.9\n",
              " 1 Choice(c0) \n",
              " 1 0 Leaf  0.5\n",
              " 1 1 Leaf  0.1\n"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "p_S_given_C"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}